// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/expression/decoded_args.h>
#include <pollux/expression/vector_function.h>
#include <pollux/type/tz/time_zone_map.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        std::optional<Timestamp> makeTimeStampFromDecodedArgs(
            vector_size_t row,
            DecodedVector *yearVector,
            DecodedVector *monthVector,
            DecodedVector *dayVector,
            DecodedVector *hourVector,
            DecodedVector *minuteVector,
            DecodedVector *microsVector) {
            // Check hour.
            auto hour = hourVector->value_at<int32_t>(row);
            if (hour < 0 || hour > 24) {
                return std::nullopt;
            }
            // Check minute.
            auto minute = minuteVector->value_at<int32_t>(row);
            if (minute < 0 || minute > 60) {
                return std::nullopt;
            }
            // Check microseconds.
            auto micros = microsVector->value_at<int64_t>(row);
            if (micros < 0) {
                return std::nullopt;
            }
            auto seconds = micros / util::kMicrosPerSec;
            if (seconds > 60 || (seconds == 60 && micros % util::kMicrosPerSec != 0)) {
                return std::nullopt;
            }

            // Year, month, day will be checked in utils::daysSinceEpochFromDate.
            Expected<int64_t> daysSinceEpoch = util::daysSinceEpochFromDate(
                yearVector->value_at<int32_t>(row),
                monthVector->value_at<int32_t>(row),
                dayVector->value_at<int32_t>(row));
            if (daysSinceEpoch.hasError()) {
                POLLUX_DCHECK(daysSinceEpoch.error().isUserError());
                return std::nullopt;
            }

            // Micros has at most 8 digits (2 for seconds + 6 for microseconds),
            // thus it's safe to cast micros from int64_t to int32_t.
            auto localMicros =
                    util::fromTime(hour, minute, 0, static_cast<int32_t>(micros));
            return util::fromDatetime(daysSinceEpoch.value(), localMicros);
        }

        void setTimestampOrNull(
            int32_t row,
            std::optional<Timestamp> timestamp,
            DecodedVector *timeZoneVector,
            FlatVector<Timestamp> *result) {
            const auto timeZoneName = timeZoneVector->value_at<StringView>(row);
            const auto *timeZone = tz::locateZone(std::string_view(timeZoneName));
            if (timestamp.has_value()) {
                timestamp->toGMT(*timeZone);
                result->set(row, *timestamp);
            } else {
                result->set_null(row, true);
            }
        }

        void setTimestampOrNull(
            int32_t row,
            std::optional<Timestamp> timestamp,
            const tz::TimeZone *timeZone,
            FlatVector<Timestamp> *result) {
            if (timestamp.has_value()) {
                timestamp->toGMT(*timeZone);
                result->set(row, *timestamp);
            } else {
                result->set_null(row, true);
            }
        }

        class MakeTimestampFunction : public exec::VectorFunction {
        public:
            MakeTimestampFunction(const tz::TimeZone *sessionTimeZone)
                : sessionTimeZone_(sessionTimeZone) {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                context.ensure_writable(rows, TIMESTAMP(), result);
                auto *resultFlatVector = result->as<FlatVector<Timestamp> >();

                exec::DecodedArgs decodedArgs(rows, args, context);
                auto *year = decodedArgs.at(0);
                auto *month = decodedArgs.at(1);
                auto *day = decodedArgs.at(2);
                auto *hour = decodedArgs.at(3);
                auto *minute = decodedArgs.at(4);
                auto *micros = decodedArgs.at(5);

                if (args.size() == 7) {
                    // If the timezone argument is specified, treat the input timestamp as the
                    // time in that timezone.
                    if (args[6]->is_constant_encoding()) {
                        auto tz =
                                args[6]->as_unchecked<ConstantVector<StringView> >()->value_at(0);
                        const tz::TimeZone *constantTimeZone = nullptr;
                        try {
                            constantTimeZone = tz::locateZone(std::string_view(tz));
                        } catch (const PolluxException &) {
                            context.setErrors(rows, std::current_exception());
                            return;
                        }
                        rows.applyToSelected([&](vector_size_t row) {
                            auto timestamp = makeTimeStampFromDecodedArgs(
                                row, year, month, day, hour, minute, micros);
                            setTimestampOrNull(
                                row, timestamp, constantTimeZone, resultFlatVector);
                        });
                    } else {
                        auto *timeZone = decodedArgs.at(6);
                        context.applyToSelectedNoThrow(rows, [&](auto row) {
                            auto timestamp = makeTimeStampFromDecodedArgs(
                                row, year, month, day, hour, minute, micros);
                            setTimestampOrNull(row, timestamp, timeZone, resultFlatVector);
                        });
                    }
                } else {
                    // Otherwise use session timezone.
                    rows.applyToSelected([&](vector_size_t row) {
                        auto timestamp = makeTimeStampFromDecodedArgs(
                            row, year, month, day, hour, minute, micros);
                        setTimestampOrNull(row, timestamp, sessionTimeZone_, resultFlatVector);
                    });
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                return {
                    exec::FunctionSignatureBuilder()
                    // precision <= 18.
                    .integerVariable("precision", "min(precision, 18)")
                    .returnType("timestamp")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("decimal(precision, 6)")
                    .build(),
                    exec::FunctionSignatureBuilder()
                    // precision <= 18.
                    .integerVariable("precision", "min(precision, 18)")
                    .returnType("timestamp")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("integer")
                    .argumentType("decimal(precision, 6)")
                    .argumentType("varchar")
                    .build(),
                };
            }

        private:
            const tz::TimeZone *sessionTimeZone_;
        };

        std::shared_ptr<exec::VectorFunction> createMakeTimestampFunction(
            const std::string & /* name */,
            const std::vector<exec::VectorFunctionArg> &inputArgs,
            const core::QueryConfig &config) {
            const auto sessionTzName = config.sessionTimezone();
            POLLUX_USER_CHECK(
                !sessionTzName.empty(),
                "make_timestamp requires session time zone to be set.");
            const auto *sessionTimeZone = tz::locateZone(sessionTzName);

            const auto &secondsType = inputArgs[5].type;
            POLLUX_USER_CHECK(
                secondsType->isShortDecimal(),
                "Seconds must be short decimal type but got {}",
                secondsType->toString());
            auto secondsScale = secondsType->asShortDecimal().scale();
            POLLUX_USER_CHECK_EQ(
                secondsScale,
                6,
                "Seconds fraction must have 6 digits for microseconds but got {}",
                secondsScale);

            return std::make_shared<MakeTimestampFunction>(sessionTimeZone);
        }
    } // namespace

    POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION(
        udf_make_timestamp,
        MakeTimestampFunction::signatures(),
        createMakeTimestampFunction);
} // namespace kumo::pollux::functions::sparksql
